import pandas as pd
import matplotlib.pyplot as plt
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
from statsmodels.tsa.stattools import adfuller

# 数据处理函数（沿用之前的）
def process_data():
    excel_file = pd.ExcelFile("ori_message.xlsx")
    years = ['2016', '2017', '2018', '2019', '2020', '2021']
    all_monthly_data = []
    for year in years:
        df = excel_file.parse(year)
        df['年'] = df['年'].ffill().astype(int)
        df['月'] = df['月'].ffill().astype(int)
        grouped = df.groupby(['年', '月']).agg({
            '流量(m3/s)': 'mean',
            '含沙量(kg/m3) ': 'mean'
        }).reset_index()
        grouped['排沙量(kg/s)'] = grouped['流量(m3/s)'] * grouped['含沙量(kg/m3) ']
        grouped = grouped.round({
            '流量(m3/s)': 2,
            '含沙量(kg/m3) ': 4,
            '排沙量(kg/s)': 2
        })
        all_monthly_data.append(grouped)
    result = pd.concat(all_monthly_data, ignore_index=True)
    result = result.sort_values(by=['年', '月']).reset_index(drop=True)
    result = result.rename(columns={'含沙量(kg/m3) ': '含沙量(kg/m3)'})
    result['时间'] = pd.to_datetime(result['年'].astype(str) + '-' + result['月'].astype(str) + '-01')
    return result

if __name__ == "__main__":
    processed_data = process_data()
    processed_data.set_index('时间', inplace=True)  # 设置时间为索引
    flow_series = processed_data['流量(m3/s)']

    # 一阶差分
    diff1 = flow_series.diff(1).dropna()
    # 二阶差分
    diff2 = diff1.diff(1).dropna()

    # 绘制一阶差分的ACF和PACF（图23）
    plt.figure(figsize=(12, 6))
    plt.suptitle('图23.一阶差分ACF、PACF图', fontsize=12)
    plt.subplot(1, 2, 1)
    plot_acf(diff1, lags=30, ax=plt.gca())
    plt.title('Autocorrelation')
    plt.subplot(1, 2, 2)
    plot_pacf(diff1, lags=30, ax=plt.gca())
    plt.title('Partial Autocorrelation')
    plt.tight_layout()
    plt.savefig('图23.一阶差分ACF、PACF图.png', dpi=300, bbox_inches='tight')
    plt.show()

    # 绘制二阶差分的ACF和PACF（图24）
    plt.figure(figsize=(12, 6))
    plt.suptitle('图24.二阶差分ACF、PACF图', fontsize=12)
    plt.subplot(1, 2, 1)
    plot_acf(diff2, lags=30, ax=plt.gca())
    plt.title('Autocorrelation')
    plt.subplot(1, 2, 2)
    plot_pacf(diff2, lags=30, ax=plt.gca())
    plt.title('Partial Autocorrelation')
    plt.tight_layout()
    plt.savefig('图24.二阶差分ACF、PACF图.png', dpi=300, bbox_inches='tight')
    plt.show()